import fla.layers.gated_deltanet as gdnp
from fla.utils import FLA_DUMP_DIR

import argparse
import itertools
import torch
import os
import glob

# CUSTOMIZE ME
HIDDEN_SIZE_L = [1024 * i for i in range(1, 9)]
EXPAND_V_L = [i for i in range(1, 3)]
HEAD_DIM_L = [i * 64 for i in range(1, 5)]
NUM_HEADS_L = [1, 2, 4, 8]
NUM_V_HEADS_L = [i for i in NUM_HEADS_L] + [NUM_HEADS_L[-1] * 2]
MODE_L = ['chunk']

LISTS = {
    "HIDDEN_SIZE_L": HIDDEN_SIZE_L, 
    "EXPAND_V_L": EXPAND_V_L, 
    "HEAD_DIM_L": HEAD_DIM_L, 
    "NUM_HEADS_L": NUM_HEADS_L, 
    "NUM_V_HEADS_L": NUM_V_HEADS_L, 
    "MODE_L": MODE_L
}

VAR_NAMES = ('hidden_size', 'expand_v', 'head_dim', 'num_heads', 'num_v_heads', 'mode')

def remove_prev_output():
    if not os.path.exists(FLA_DUMP_DIR):
        os.makedirs(FLA_DUMP_DIR)
        print(f"Directory created: {FLA_DUMP_DIR}")
        return

    if os.path.isdir(FLA_DUMP_DIR):
        # Get all files in the directory (not subdirectories)
        files = glob.glob(os.path.join(FLA_DUMP_DIR, "*"))
        
        # Remove each file
        for file_path in files:
            if os.path.isfile(file_path):
                os.remove(file_path)
                print(f"Removed file: {file_path}")
        
        print(f"All files removed from: {FLA_DUMP_DIR}")
        return 
    print(f"Path exists but is not a directory: {FLA_DUMP_DIR}")

def format_iter_info(x):
    s = "{\n"
    for i in range(6):
        s += f"\t{VAR_NAMES[i]}: {x[i]}\n"
    s += "}"
    return s

def normalize(x):
    x = (int(x[0]), int(x[1]), int(x[2]), int(x[3]), int(x[4]), x[5])
    return x

# currencty allow num_vheads be only up to 2x of num_heads
def legal(x):
    return x[3] == x[4] or x[3] * 2 == x[4]

def create_cuseqs(seq_lens:int):
    """精确创建8192长度的序列"""
    # 随机决定子句数量
    m = torch.randint(3, 6, (1,)).item()  # 10-20个子句
    # 生成随机子句长度
    clause_lengths = torch.randint(seq_lens // (2 * m), seq_lens // m, (m-1,))  # 前m-1个子句
    last_clause = seq_lens - clause_lengths.sum().item()
    clause_lengths = torch.cat([clause_lengths, torch.tensor([last_clause])])
    # 创建cu_seqs
    cu_seqs = torch.cat([
        torch.tensor([0]), 
        torch.cumsum(clause_lengths, dim=0)
    ])
    print(f"子句数量: {m}")
    print(f"每个子句长度: {clause_lengths.tolist()}")
    # print(f"累计序列位置 (cu_seqs): {cu_seqs.tolist()}")
    # print(f"总序列长度: {cu_seqs[-1].item()}")
    return cu_seqs.to('npu:0')

# GDN RUN STEP
def step(x):
    assert legal(x)
    print(f"Running {format_iter_info(x)}")

    gdn = gdnp.GatedDeltaNet(
        hidden_size = x[0],
        expand_v = float(x[1]),
        head_dim = x[2],
        num_heads = x[3],
        num_v_heads = x[4],
        mode = x[5]
    )
    print(f"Successfully instantiated and initialized GDN for given params")

    # hidden_states is basically a 3d tensor with shape (batch_size, seq_len, hidden_size)
    # FIXME TODO investigate different shapes
    # hidden_states = torch.randn(8, 128, gdn.hidden_size, dtype=torch.float16)
    seq_len = 8 * 1024
    hidden_states = torch.randn(1, seq_len, gdn.hidden_size, dtype=torch.float16)
    print(f"Successfully initialized hidden_states with shape {hidden_states.shape} on {torch.get_default_device()}")
    print(f"Forwarding...")
    # FIXME TODO collect info and compare with interpreter
    # FIXME TODO adjust kernels if support modes other than chunk
    # INFO chunk case: forward -> chunk_gated_delta_rule -> ChunkGatedDeltaRuleFunction.apply -> {
    #   chunk_gated_delta_rule_fwd
    #   chunk_gated_delta_rule_bwd
    # }
    # FIXME TODO !!! replace called kernels with adjusted versions from kernel_tests
    cu_seqlens = create_cuseqs(hidden_states.shape[1])
    kwargs = {'cu_seqlens': cu_seqlens,}

    output, _, _ = gdn.forward(hidden_states, **kwargs)
    print(f"Backwarding...")
    (output-1).sum().backward()

    print(f"Finished")

    print()

def main():
    remove_prev_output()    
    
    torch.manual_seed(42)
    torch.set_default_device('npu')
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--all', action='store_true')
    group.add_argument('--iteration', type=str, nargs='+', metavar=VAR_NAMES)
    
    args = parser.parse_args()
    A, B, C, D, E, F = LISTS.values()

    def print_values():
        print('Available values:')
        for n, l in LISTS.items():
            print(f"\t{n}: {l}")

    def parse_range(range_str, allowed_list):
        if '..' in range_str:
            start, end = range_str.split('..', 1)
            try:
                start_conv = type(allowed_list[0])(start) if allowed_list else start
                end_conv = type(allowed_list[0])(end) if allowed_list else end
            except (ValueError, TypeError):
                return None
                
            if start_conv not in allowed_list or end_conv not in allowed_list:
                return None
                
            start_idx = allowed_list.index(start_conv)
            end_idx = allowed_list.index(end_conv)
            if start_idx <= end_idx:
                return allowed_list[start_idx:end_idx + 1]
            else:
                return allowed_list[end_idx:start_idx + 1][::-1]
        else:
            try:
                value_conv = type(allowed_list[0])(range_str) if allowed_list else range_str
                return [value_conv] if value_conv in allowed_list else None
            except (ValueError, TypeError):
                return None

    if args.all:
        for x in itertools.product(A, B, C, D, E, F):
            x = normalize(x)
            if legal(x):
                step(x)
    
    elif args.iteration:
        if len(args.iteration) != 6:
            print(f"Error: Expected 6 arguments, got {len(args.iteration)}")
            print(f"Usage: --iteration {' '.join(VAR_NAMES)}")
            print("You can use ranges: value..value")
            print_values()
            exit(1)

        parsed_args = []
        for i, arg in enumerate(args.iteration):
            try:
                allowed_list = [A, B, C, D, E, F][i]
                if '..' in arg:
                    values = parse_range(arg, allowed_list)
                    if values is None:
                        raise ValueError()
                    parsed_args.append(values)
                else:
                    value_conv = type(allowed_list[0])(arg) if allowed_list else arg
                    if value_conv not in allowed_list:
                        raise ValueError()
                    parsed_args.append([value_conv])
            except (ValueError, TypeError):
                print(f"Error: Invalid value '{arg}' for {VAR_NAMES[i]}")
                print(f"Available values for {VAR_NAMES[i]}: {allowed_list}")
                exit(1)

        for combination in itertools.product(*parsed_args):
            x = normalize(combination)
            if legal(x):
                step(x)
    
    else:
        print_values()
        print(f'Specify --all or --iteration {" ".join(VAR_NAMES)}')
        print("You can use ranges: v1..v2")

if __name__ == '__main__':
    main()
